package repo

import (
	"context"
	"fmt"

	"github.com/go-git/go-git/v5"
	"github.com/go-git/go-git/v5/plumbing"
	"github.com/go-git/go-git/v5/plumbing/transport/http"
	gitssh "github.com/go-git/go-git/v5/plumbing/transport/ssh"
	"golang.org/x/crypto/ssh"

	"github.com/dstackai/dstack/runner/internal/log"
)

type Manager struct {
	ctx       context.Context
	localPath string
	clo       git.CloneOptions
	branch    string
	hash      string
}

func NewManager(ctx context.Context, url, branch, hash string, singleBranch bool) *Manager {
	ctx = log.AppendArgsCtx(ctx, "url", url, "branch", branch, "hash", hash)
	m := &Manager{
		ctx:    ctx,
		branch: branch,
		hash:   hash,
		clo: git.CloneOptions{
			URL:               url,
			RecurseSubmodules: git.DefaultSubmoduleRecursionDepth,
			SingleBranch:      singleBranch,
		},
	}
	// Only set ReferenceName if branch is non-empty
	// If empty, it will default to HEAD in CloneOptions.Validate()
	if branch != "" {
		m.clo.ReferenceName = plumbing.NewBranchReferenceName(branch)
	}

	return m
}

func (m *Manager) WithLocalPath(path string) *Manager {
	m.localPath = path
	m.ctx = log.AppendArgsCtx(m.ctx, "path", path)
	return m
}

// TODO: works with Github, possibly not with others
func (m *Manager) WithTokenAuth(token string) *Manager {
	auth := &http.BasicAuth{
		Username: "anything",
		Password: token,
	}
	m.clo.Auth = auth
	return m
}

func (m *Manager) WithSSHAuth(pem, password string) *Manager {
	keys, err := gitssh.NewPublicKeys("git", []byte(pem), password)
	if err != nil {
		log.Warning(m.ctx, "fail to parse SSH private key", "err", err)
	} else {
		keys.HostKeyCallback = ssh.InsecureIgnoreHostKey()
		m.clo.Auth = keys
	}
	return m
}

func (m *Manager) Checkout() error {
	log.Info(m.ctx, "git checkout", "auth", fmt.Sprintf("%T", (&m.clo).Auth))
	ref, err := git.PlainClone(m.localPath, false, &m.clo)
	if err != nil {
		return fmt.Errorf("clone repo: %w", err)
	}
	if ref != nil {
		var cho git.CheckoutOptions
		needCheckout := false

		if m.branch != "" {
			branchRef, err := ref.Reference(m.clo.ReferenceName, true)
			if err != nil {
				return fmt.Errorf("get branch reference: %w", err)
			}
			if m.hash == "" || m.hash == branchRef.Hash().String() {
				// Hash is empty or matches branch head: checkout branch
				cho.Branch = m.clo.ReferenceName
				needCheckout = true
			} else {
				// Hash is specified and different: checkout by hash
				cho.Hash = plumbing.NewHash(m.hash)
				needCheckout = true
			}
		} else {
			// Branch is empty: checkout by hash if specified, otherwise HEAD is already checked out
			if m.hash != "" {
				cho.Hash = plumbing.NewHash(m.hash)
				needCheckout = true
			}
			// If hash is also empty, HEAD is already checked out by clone, no need to checkout again
		}

		if needCheckout {
			workTree, err := ref.Worktree()
			if err != nil {
				return fmt.Errorf("get worktree: %w", err)
			}
			err = workTree.Checkout(&cho)
			if err != nil {
				return fmt.Errorf("checkout: %w", err)
			}
		}
	} else {
		log.Warning(m.ctx, "git clone ref==nil")
	}

	return nil
}

func (m *Manager) URL() string {
	return m.clo.URL
}

func (m *Manager) SetConfig(name, email string) error {
	repo, err := git.PlainOpen(m.localPath)
	if err != nil {
		return fmt.Errorf("open repo: %w", err)
	}
	config, err := repo.Config()
	if err != nil {
		return fmt.Errorf("get repo config: %w", err)
	}
	config.User.Name = name
	config.User.Email = email
	if err := repo.SetConfig(config); err != nil {
		return fmt.Errorf("set repo config: %w", err)
	}
	return nil
}
